// Copyright(c) 2016 Ruoyu Fan (Windy Darian), Xueyin Wan
// MIT License.

#include "Context.h"

#include <GLFW/glfw3.h>

#include <unordered_set>
#include <iostream>
#include <cstring>
#include <set>
#include <stdexcept>

#ifdef NDEBUG
// if not debugging
const bool ENABLE_VALIDATION_LAYERS = false;
#else
const bool ENABLE_VALIDATION_LAYERS = true;
#endif

const std::vector<const char*> VALIDATION_LAYERS = {
	"VK_LAYER_LUNARG_standard_validation"
};

const std::vector<const char*> DEVICE_EXTENSIONS = {
	VK_KHR_SWAPCHAIN_EXTENSION_NAME
};

VContext::VContext(GLFWwindow* window)
{
	if (!window)
	{
		throw std::runtime_error("invalid window");
	}

	this->window = window;

	initVulkan();
}

VContext::~VContext() = default;

std::pair<int, int> VContext::getWindowFrameBufferSize() const
{
	int framebuffer_width, framebuffer_height;
	glfwGetFramebufferSize(window, &framebuffer_width, &framebuffer_height);
	return std::make_pair(framebuffer_width, framebuffer_height);
}

// the debug callback function that Vulkan runs
VkBool32 debugCallback(
		VkDebugReportFlagsEXT flags,
		VkDebugReportObjectTypeEXT objType,
		uint64_t obj,
		size_t location,
		int32_t code,
		const char* layerPrefix,
		const char* msg,
		void* userData)
{

	std::cerr << "validation layer: " << msg << std::endl;
	throw std::runtime_error(msg);
	return VK_FALSE;
}

QueueFamilyIndices QueueFamilyIndices::findQueueFamilies(vk::PhysicalDevice device, vk::SurfaceKHR surface)

{
	QueueFamilyIndices indices;
	auto queuefamilies = device.getQueueFamilyProperties();

	int i = 0;
	for (const auto& queuefamily : queuefamilies)
	{
		if (queuefamily.queueCount > 0 && queuefamily.queueFlags & vk::QueueFlagBits::eGraphics)
		{
			indices.graphics_family = i;
		}

		if (queuefamily.queueCount > 0 && queuefamily.queueFlags & vk::QueueFlagBits::eCompute)
		{
			indices.compute_family = i;
		}

		if (queuefamily.queueCount > 0 && device.getSurfaceSupportKHR(i,surface))
		{
			indices.present_family = i;
		}

		if (indices.isComplete()) {
			break;
		}

		i++;
	}

	return indices;
}

std::vector<const char*> getRequiredExtensions()
{
	std::vector<const char*> extensions;

	unsigned int glfwExtensionCount = 0;
	const char** glfwExtensions;
	glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);

	for (unsigned int i = 0; i < glfwExtensionCount; i++)
	{
		extensions.push_back(glfwExtensions[i]);
	}

	if (ENABLE_VALIDATION_LAYERS)
	{
		extensions.push_back(VK_EXT_DEBUG_REPORT_EXTENSION_NAME);
	}

	return extensions;
	// should I free sth after?
}

VkResult VContext::CreateDebugReportCallbackEXT(VkInstance instance, const VkDebugReportCallbackCreateInfoEXT* pCreateInfo, const VkAllocationCallbacks* pAllocator, VkDebugReportCallbackEXT* pCallback)
{
	auto func = (PFN_vkCreateDebugReportCallbackEXT)vkGetInstanceProcAddr(instance, "vkCreateDebugReportCallbackEXT");
	if (func != nullptr)
	{
		return func(instance, pCreateInfo, pAllocator, pCallback);
	}
	else
	{
		return VK_ERROR_EXTENSION_NOT_PRESENT;
	}
}

void VContext::DestroyDebugReportCallbackEXT(VkInstance instance, VkDebugReportCallbackEXT callback, const VkAllocationCallbacks* pAllocator)
{
	auto func = (PFN_vkDestroyDebugReportCallbackEXT)vkGetInstanceProcAddr(instance, "vkDestroyDebugReportCallbackEXT");
	if (func != nullptr) {
		func(instance, callback, pAllocator);
	}
}

bool checkValidationLayerSupport()
{
	auto availableLayers = vk::enumerateInstanceLayerProperties();

	for (const char* layerName : VALIDATION_LAYERS)
	{
		bool layerFound = false;

		for (const auto& layerProperties : availableLayers)
		{
			if (strcmp(layerName, layerProperties.layerName) == 0)
			{
				layerFound = true;
				break;
			}
		}

		if (!layerFound) {
			return false;
		}
	}

	return true;
}

void VContext::createInstance()
{
	if (ENABLE_VALIDATION_LAYERS && !checkValidationLayerSupport())
	{
		throw std::runtime_error("validation layers requested, but not available!");
	}

	// Getting Vulkan instance extensions required by GLFW
	auto glfwExtensions = getRequiredExtensions();

	// Getting Vulkan supported extensions
	auto extensions = vk::enumerateInstanceExtensionProperties();
	std::unordered_set<std::string> supportedExtensionNames;
	for (const auto& extension : extensions)
	{
		supportedExtensionNames.insert(std::string(extension.extensionName));
	}

	// Print Vulkan supported extensions
	std::cout << "available extensions:" << std::endl;
	for (const auto& name : supportedExtensionNames) {
		std::cout << "\t" << name << std::endl;
	}
	// Check for and print any unsupported extension
	for (const auto& extension_name : glfwExtensions)
	{
		std::string name(extension_name);
		if (supportedExtensionNames.count(name) <= 0)
		{
			std::cout << "unsupported extension required by GLFW: " << name << std::endl;
		}
	}

    auto app_info = vk::ApplicationInfo()
            .setApiVersion(VK_API_VERSION_1_0);

    auto instance_info = vk::InstanceCreateInfo()
        .setPApplicationInfo(&app_info)
        .setEnabledExtensionCount(static_cast<uint32_t>(glfwExtensions.size()))
        .setPpEnabledExtensionNames(glfwExtensions.data());

	if (ENABLE_VALIDATION_LAYERS) {
        instance_info.setEnabledLayerCount(static_cast<uint32_t>(VALIDATION_LAYERS.size()))
                .setPpEnabledLayerNames(VALIDATION_LAYERS.data());
	}
	else {
        instance_info.setEnabledLayerCount(0);
	}

	instance = VRaii<vk::Instance>{
		vk::createInstance(instance_info, nullptr),
		[](vk::Instance& obj)
		{
			obj.destroy();
		}
	};

}

void VContext::setupDebugCallback()
{
	if (!ENABLE_VALIDATION_LAYERS) return;

	VkDebugReportCallbackCreateInfoEXT createInfo = {};
	createInfo.sType = VK_STRUCTURE_TYPE_DEBUG_REPORT_CALLBACK_CREATE_INFO_EXT;
	createInfo.flags = VK_DEBUG_REPORT_ERROR_BIT_EXT | VK_DEBUG_REPORT_WARNING_BIT_EXT;
	createInfo.pfnCallback = (PFN_vkDebugReportCallbackEXT)debugCallback;

	VkDebugReportCallbackEXT temp_callback;
	if (CreateDebugReportCallbackEXT(instance.get(), &createInfo, nullptr, &temp_callback) != VK_SUCCESS)
	{
		throw std::runtime_error("failed to set up debug callback!");
	}

    auto deleter = [this](VkDebugReportCallbackEXT& obj)
    {
        DestroyDebugReportCallbackEXT(instance.get(), obj, nullptr);
    };

	callback = VRaii<VkDebugReportCallbackEXT>{
		temp_callback,deleter
	};

}

bool checkDeviceExtensionSupport(vk::PhysicalDevice device)
{
    auto available_extensions = device.enumerateDeviceExtensionProperties();

	std::set<std::string> required_extensions(DEVICE_EXTENSIONS.begin(), DEVICE_EXTENSIONS.end());

	for (const auto& extension : available_extensions)
	{
		required_extensions.erase(extension.extensionName);
	}

	return required_extensions.empty();
}


bool isDeviceSuitable(vk::PhysicalDevice device, vk::SurfaceKHR window_surface)
{
	QueueFamilyIndices indices = QueueFamilyIndices::findQueueFamilies(device, window_surface);

	bool extensions_supported = checkDeviceExtensionSupport(device);

	bool swap_chain_adequate = false;
	if (extensions_supported)
	{
		auto swap_chain_support = SwapChainSupportDetails::querySwapChainSupport(device, window_surface);
		swap_chain_adequate = !swap_chain_support.formats.empty() && !swap_chain_support.present_modes.empty();
	}

	return indices.isComplete() && extensions_supported && swap_chain_adequate;
}

// Pick up a graphics card to use
void VContext::pickPhysicalDevice()
{
	// This object will be implicitly destroyed when the vk::Instance is destroyed, so we don't need to add a delete wrapper.
	vk::PhysicalDevice physial_device;
    auto devices = instance.get().enumeratePhysicalDevices();

	for (const auto& device : devices)
	{
		if (isDeviceSuitable(device, window_surface.get()))
		{
			physial_device = device;
			break;
		}
	}

	if (!physial_device)
	{
		throw std::runtime_error("Failed to find a suitable GPU!");
	}
	else
	{
        auto properties = physial_device.getProperties();
		std::cout << "Current Device: " << properties.deviceName << std::endl;
	}

	this->physical_device = physial_device;

	this->physical_device_properties = physical_device.getProperties();
}

void VContext::findQueueFamilyIndices()
{
	queue_family_indices = QueueFamilyIndices::findQueueFamilies(physical_device, window_surface);

	if (!queue_family_indices.isComplete())
	{
		throw std::runtime_error("Queue family indices not complete!");
	}
}

// Needs to be called right after instance creation because it may influence device selection
void VContext::createWindowSurface()
{
	VkSurfaceKHR temp_surface;
	auto result = glfwCreateWindowSurface(instance.get(), window, nullptr, &temp_surface);

	if (result != VK_SUCCESS)
	{
		throw std::runtime_error("Failed to create window surface!");
	}

	window_surface = VRaii<vk::SurfaceKHR>{
		temp_surface,
		[this](vk::SurfaceKHR& surface)
		{
			instance.get().destroySurfaceKHR(surface);
		}
	};
}

void VContext::createLogicalDevice()
{
	QueueFamilyIndices indices = QueueFamilyIndices::findQueueFamilies(physical_device, window_surface);

	std::vector <vk::DeviceQueueCreateInfo> queue_create_infos;
	std::set<int> queue_families = { indices.graphics_family, indices.present_family, indices.compute_family };

	float queue_priority = 1.0f;
	for (int family : queue_families)
	{
        // Create a graphics queue
        auto queue_create_info = vk::DeviceQueueCreateInfo()
            .setQueueCount(1)
            .setQueueFamilyIndex(indices.graphics_family)
            .setPQueuePriorities(&queue_priority);

		queue_create_infos.push_back(queue_create_info);
	}
	
	// Specify used device features
	vk::PhysicalDeviceFeatures device_features = {}; // Everything is by default VK_FALSE
    // Create the logical device
    auto device_create_info = vk::DeviceCreateInfo()
        .setQueueCreateInfoCount(static_cast<uint32_t> (queue_create_infos.size()))
        .setPQueueCreateInfos(queue_create_infos.data())
        .setPEnabledFeatures(&device_features)
        .setEnabledExtensionCount(static_cast<uint32_t>(DEVICE_EXTENSIONS.size()))
        .setPpEnabledExtensionNames(DEVICE_EXTENSIONS.data());

	if (ENABLE_VALIDATION_LAYERS)
	{
        device_create_info.setEnabledLayerCount(static_cast<uint32_t>(VALIDATION_LAYERS.size()))
                .setPpEnabledLayerNames(VALIDATION_LAYERS.data());
	}
	else
	{
        device_create_info.setEnabledLayerCount(0);
	}
    auto temp_device = physical_device.createDevice(device_create_info);
	graphics_device = VRaii<vk::Device>{
		temp_device,
		[](vk::Device& device)
		{
			device.destroy();
		}
	};
	auto device = graphics_device.get();

	graphics_queue = device.getQueue(indices.graphics_family, 0);
	present_queue = device.getQueue(indices.present_family, 0);
	compute_queue = device.getQueue(indices.compute_family, 0);
}

void VContext::createCommandPools()
{
	auto device = graphics_device.get();
	auto raii_commandpool_deleter = [device](vk::CommandPool& obj)
	{
		device.destroyCommandPool(obj);
	};

	auto& indices = queue_family_indices;

	// graphics_queue_command_pool
	{
        auto pool_info = vk::CommandPoolCreateInfo(vk::CommandPoolCreateFlags(),indices.graphics_family);
		graphics_queue_command_pool = VRaii<vk::CommandPool>(
			device.createCommandPool(pool_info, nullptr),
			raii_commandpool_deleter
		);
	}
	
	// compute_queue_command_pool
	{
        auto cmd_pool_info = vk::CommandPoolCreateInfo(vk::CommandPoolCreateFlagBits::eResetCommandBuffer,indices.graphics_family);
		compute_queue_command_pool = VRaii<vk::CommandPool>(
			device.createCommandPool(cmd_pool_info, nullptr),
			raii_commandpool_deleter
		);
	}

}

SwapChainSupportDetails SwapChainSupportDetails::querySwapChainSupport(vk::PhysicalDevice device, vk::SurfaceKHR surface)
{
	SwapChainSupportDetails details;
    details.capabilities = device.getSurfaceCapabilitiesKHR(surface);

	// Getting supported surface formats
    details.formats = device.getSurfaceFormatsKHR(surface);

	// Getting supported present modes
    details.present_modes = device.getSurfacePresentModesKHR(surface);

	return details;
}